package com.onion.utils.http;

import java.io.IOException;
import java.io.InputStream;

/**文件下载限流工具
 * @author liujiankuan
 */
public class FileDownloadLimitUtil {
    /**
     * KB代表的字节数
     */
    private static final Long KB = 1024L;

    /**
     * 获取限流后的inputstream
     * @param inputStream
     * @param maxRate 限流速率单位KB/s
     * @return 限流后的inputstream
     */
    public static InputStream getLimitInputStream(InputStream inputStream, int maxRate){
        if (inputStream == null){
            throw new RuntimeException("io流为空");
        }
        BandwidthLimiter bandwidthLimiter = null;
        if (maxRate > 0){
            bandwidthLimiter = new BandwidthLimiter(maxRate);
        }
        return new LimitInputStream(inputStream, bandwidthLimiter);
    }

    private static class BandwidthLimiter {
        /**
         * 已经发送/读取的字节数
         */
        private int bytesWillBeSentOrReceive = 0;
        /**
         * 上一次接收到字节流的时间戳——单位纳秒
         */
        private long lastPieceSentOrReceiveTick = System.nanoTime();
        /**
         * 允许的最大速率，默认为 1024KB/s
         */
        private int maxRate = 1024;

        /**
         * 一个chunk的大小，单位byte。设置一个块的大小为1M
         */
        private  Long CHUNK_LENGTH;
        /**
         * 在maxRate的速率下，通过chunk大小的字节流要多少时间(纳秒)
         */
        private long timeCostPerChunk;

        public BandwidthLimiter(int maxRate) {
            this.CHUNK_LENGTH = 1024L*maxRate;
            this.setMaxRate(maxRate);
        }

        /**
         * 动态调整最大速率
         * @param maxRate
         */
        public void setMaxRate(int maxRate) {
            if (maxRate < 0) {
                throw new IllegalArgumentException("maxRate can not less than 0");
            }
            this.maxRate = maxRate;
            if (maxRate == 0) {
                this.timeCostPerChunk = 0;
            } else {
                this.timeCostPerChunk = (1000000000L * CHUNK_LENGTH) / (this.maxRate * KB);
            }

        }
        public void limitNextBytes() {
            this.limitNextBytes(1);
        }

        public void limitNextBytes(int len) {
            this.bytesWillBeSentOrReceive += len;
            long round = this.bytesWillBeSentOrReceive / CHUNK_LENGTH;
            if (round == 0){
                return;
            }
            long nowTick = System.nanoTime();
            long passTime = nowTick - this.lastPieceSentOrReceiveTick;
            long missedTime = this.timeCostPerChunk * round - passTime;
            if (missedTime > 0) {
                try {
                    Thread.sleep(missedTime / 1000000, (int) (missedTime % 1000000));
                } catch (InterruptedException e) {
                    System.out.println(e);
                }
            }
            this.bytesWillBeSentOrReceive -= CHUNK_LENGTH*round;
            this.lastPieceSentOrReceiveTick = nowTick + (missedTime > 0 ? missedTime : 0);
        }
    }

    private static class LimitInputStream extends InputStream {

        private InputStream inputStream;

        private BandwidthLimiter bandwidthLimiter;

        public LimitInputStream(InputStream inputStream, BandwidthLimiter bandwidthLimiter) {
            this.inputStream = inputStream;
            this.bandwidthLimiter = bandwidthLimiter;
        }

        @Override
        public int read() throws IOException {
            if (bandwidthLimiter != null) {
                bandwidthLimiter.limitNextBytes();
            }
            return inputStream.read();

        }
        @Override
        public int read(byte[] b, int off, int len) throws IOException {
            if (bandwidthLimiter != null) {
                bandwidthLimiter.limitNextBytes(len);
            }
            return inputStream.read(b, off, len);
        }

        @Override
        public int read(byte[] b) throws IOException {
            if (bandwidthLimiter != null && b.length > 0) {
                bandwidthLimiter.limitNextBytes(b.length);
            }
            return inputStream.read(b);
        }
    }

}